import unittest

import tensorflow as tf

from model.tf.plainnet import PlainNet

HEIGHT, WIDTH, DEPTH = 32, 32, 3


class MyTestCase(unittest.TestCase):
    def setUp(self) -> None:
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                tf.config.experimental.set_virtual_device_configuration(
                    gpus[0],
                    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
            except RuntimeError as e:
                print(e)

    def test_plain_net(self):
        plain_net = PlainNet()
        plain_net(tf.keras.layers.Input([HEIGHT, WIDTH, DEPTH]))
        print(plain_net.name)
        print(plain_net.summary())
        self.assertEqual((40, 200), plain_net.compute_output_shape([40, HEIGHT, WIDTH, DEPTH]))
